{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inverted index, queries and TF-IDF by hand in Python\n",
    "In this notebook, I provide functions for cleaning, indexing and executing word and phrase queries, and show how to compute TF-IDF scores. All of this is illustrated on a toy collection of documents. We need very few external librairies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import nltk\n",
    "import operator\n",
    "from nltk import word_tokenize\n",
    "stemmer = nltk.stem.PorterStemmer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleaning and indexing\n",
    "The inverted index contains, for each unique term, the position(s) of the term in each document of the collection. It can be structured as a dictionary of dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def clean_tokenize(doc):\n",
    "    words = word_tokenize(doc.lower())\n",
    "    return words\n",
    "\n",
    "def index_one_doc(doc,to_stem):\n",
    "    '''\n",
    "    creates dict (tok,positions) for each tok in the document (as term list)\n",
    "    '''\n",
    "    tokpos = dict()\n",
    "    for t_idx,tok in enumerate(doc):\n",
    "        if to_stem:\n",
    "           tok = stemmer.stem(tok)\n",
    "        if tok in tokpos:\n",
    "            tokpos[tok].append(t_idx)\n",
    "        else:\n",
    "            tokpos[tok] = [t_idx]\n",
    "    return tokpos"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what we get when indexing our toy corpus. We first clean the documents and inspect the positional lists for individual documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'an': [0], 'engin': [1, 5], 'may': [2], 'design': [3], 'car': [4], 'of': [6], 'all': [7], 'sort': [8]}\n",
      "{'an': [0], 'engineer': [1], 'may': [2], 'design': [3], 'car': [4], 'engines': [5], 'of': [6], 'all': [7], 'sorts': [8]}\n"
     ]
    }
   ],
   "source": [
    "docs = ['The quick brown fox jumps over the lazy dog',\n",
    "        'The brown quick fox jumps over the lazy dog',\n",
    "        'Luke is the mechanical and electrical engineer of the group',\n",
    "        'Instead of buying a new engine, buy a new car',\n",
    "        'An engineer may design car engines of all sorts',\n",
    "        'Engineers use logic, but not necessarily imagination',\n",
    "        'Logic will take you from A to Z, imagination will take you everywhere.',\n",
    "        'Continuous effort, not strength or intelligence, is the key to \\\n",
    "        unlocking our potential. And curiosity.',\n",
    "        'It’s OK to have your eggs in one basket as long as you control what \\\n",
    "        happens to that basket.'\n",
    "        ]\n",
    "\n",
    "cleaned_docs = []\n",
    "for doc in docs:\n",
    "    to_app = clean_tokenize(doc)\n",
    "    cleaned_docs.append(to_app)\n",
    "\n",
    "print(index_one_doc(cleaned_docs[4],to_stem=True))\n",
    "print(index_one_doc(cleaned_docs[4],to_stem=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that obviously, stemming affects our inverted index. In the example above, 'engineer' and 'engines' are collapsed to the same 'engin' term when stemming is used. Let's now construct the full inverted index. Note that we make the following choices: our (1) queries will not be case sensitive, (2) we are indexing punctuation marks, (3) we are indexing stopwords and thus can have more expressive queries (that capture negation, for instance). We also construct two inverted indexes: one using stemming and one not using stemming."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'the': {0: [0, 6], 1: [0, 6], 2: [2, 8], 7: [9]}, 'quick': {0: [1], 1: [2]}, 'brown': {0: [2], 1: [1]}, 'fox': {0: [3], 1: [3]}, 'jump': {0: [4], 1: [4]}, 'over': {0: [5], 1: [5]}, 'lazi': {0: [7], 1: [7]}, 'dog': {0: [8], 1: [8]}, 'luke': {2: [0]}, 'is': {2: [1], 7: [8]}, 'mechan': {2: [3]}, 'and': {2: [4], 7: [16]}, 'electr': {2: [5]}, 'engin': {2: [6], 3: [5], 4: [1, 5], 5: [0]}, 'of': {2: [7], 3: [1], 4: [6]}, 'group': {2: [9]}, 'instead': {3: [0]}, 'buy': {3: [2, 7]}, 'a': {3: [3, 8], 6: [5]}, 'new': {3: [4, 9]}, ',': {3: [6], 5: [3], 6: [8], 7: [2, 7]}, 'car': {3: [10], 4: [4]}, 'an': {4: [0]}, 'may': {4: [2]}, 'design': {4: [3]}, 'all': {4: [7]}, 'sort': {4: [8]}, 'use': {5: [1]}, 'logic': {5: [2], 6: [0]}, 'but': {5: [4]}, 'not': {5: [5], 7: [3]}, 'necessarili': {5: [6]}, 'imagin': {5: [7], 6: [9]}, 'will': {6: [1, 10]}, 'take': {6: [2, 11]}, 'you': {6: [3, 12], 8: [14]}, 'from': {6: [4]}, 'to': {6: [6], 7: [11], 8: [4, 18]}, 'z': {6: [7]}, 'everywher': {6: [13]}, '.': {6: [14], 7: [15, 18], 8: [21]}, 'continu': {7: [0]}, 'effort': {7: [1]}, 'strength': {7: [4]}, 'or': {7: [5]}, 'intellig': {7: [6]}, 'key': {7: [10]}, 'unlock': {7: [12]}, 'our': {7: [13]}, 'potenti': {7: [14]}, 'curios': {7: [17]}, 'it': {8: [0]}, '’': {8: [1]}, 's': {8: [2]}, 'ok': {8: [3]}, 'have': {8: [5]}, 'your': {8: [6]}, 'egg': {8: [7]}, 'in': {8: [8]}, 'one': {8: [9]}, 'basket': {8: [10, 20]}, 'as': {8: [11, 13]}, 'long': {8: [12]}, 'control': {8: [15]}, 'what': {8: [16]}, 'happen': {8: [17]}, 'that': {8: [19]}}\n"
     ]
    }
   ],
   "source": [
    "inverted_index = dict()\n",
    "inverted_index_stem = dict()\n",
    "\n",
    "for d_idx,doc in enumerate(cleaned_docs):\n",
    "    \n",
    "    poslists_s = index_one_doc(doc,to_stem=True) # get positions of each token in the doc\n",
    "    for tok,poslist_s in poslists_s.items():\n",
    "        if tok in inverted_index_stem:\n",
    "            inverted_index_stem[tok][d_idx] = poslist_s # update\n",
    "        else:\n",
    "            inverted_index_stem[tok] = dict()\n",
    "            inverted_index_stem[tok][d_idx] = poslist_s # initialize\n",
    "    \n",
    "    poslists = index_one_doc(doc,to_stem=False)\n",
    "    for tok,poslist in poslists.items():\n",
    "        if tok in inverted_index:\n",
    "            inverted_index[tok][d_idx] = poslist\n",
    "        else:\n",
    "            inverted_index[tok] = dict()\n",
    "            inverted_index[tok][d_idx] = poslist\n",
    "\n",
    "print(inverted_index_stem)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that our full inverted index contains the positional lists of each term in each document of the collection.\n",
    "\n",
    "## Querying\n",
    "Let's now define some functions to execute word and phrase queries over our collection:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def at_least_one_unigram(query,inverted_index):\n",
    "    '''\n",
    "    returns the indexes of the docs containing *at least one* query unigrams\n",
    "    the query is a list of unigrams\n",
    "    '''\n",
    "    \n",
    "    to_return = []\n",
    "    for unigram in query:\n",
    "        if unigram in inverted_index:\n",
    "            to_return.extend(list(inverted_index[unigram].keys()))\n",
    "    return list(set(to_return))\n",
    "\n",
    "def all_unigrams(query,inverted_index):\n",
    "    '''\n",
    "    returns the indexes of the docs containing *all* query unigrams\n",
    "    the query is a list of unigrams\n",
    "    '''\n",
    "    \n",
    "    to_return = []\n",
    "    for unigram in query:\n",
    "        if unigram in inverted_index:\n",
    "            to_return.append(set(list(inverted_index[unigram].keys())))\n",
    "        else:\n",
    "            to_return.append(set())\n",
    "            break\n",
    "    to_return = to_return[0].intersection(*to_return)\n",
    "    return list(to_return)\n",
    "\n",
    "def ngrams(query,inverted_index):\n",
    "    '''\n",
    "    returns the indexes of the docs containing all unigrams in same order as the query\n",
    "    the query is a list of unigrams\n",
    "    '''\n",
    "    candidate_docs = all_unigrams(query,inverted_index)   \n",
    "    \n",
    "    to_return = []\n",
    "    for doc in candidate_docs:\n",
    "        poslists = []\n",
    "        for unigram in query:\n",
    "            to_append = inverted_index[unigram][doc]\n",
    "            if isinstance(to_append, int):\n",
    "                poslists.append([to_append])\n",
    "            else:\n",
    "                poslists.append(to_append)\n",
    "        # test whether the query words are consecutive    \n",
    "        poslists_sub = [[elt-idx for elt in poslist] for idx,poslist in enumerate(poslists)]\n",
    "        if set(poslists_sub[0]).intersection(*poslists_sub):\n",
    "            to_return.append(doc)\n",
    "    return to_return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's perform a few queries to illustrate how the functions work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3, 4]\n",
      "[3]\n",
      "[2, 3, 4, 5]\n",
      "[3, 4]\n",
      "[0, 1]\n",
      "[0]\n"
     ]
    }
   ],
   "source": [
    "# 1) docs containing engine OR car\n",
    "query = ['engine','car']\n",
    "print(at_least_one_unigram(query,inverted_index))\n",
    "\n",
    "# 2) docs containing engine AND car\n",
    "print(all_unigrams(query,inverted_index))\n",
    "\n",
    "# repeat 1) and 2) above, using stemming\n",
    "query_stemmed = [stemmer.stem(elt) for elt in query]\n",
    "print(at_least_one_unigram(query_stemmed,inverted_index_stem))\n",
    "print(all_unigrams(query_stemmed,inverted_index_stem))\n",
    "\n",
    "# 3) all unigrams vs. phrase query\n",
    "query = ['quick','brown']\n",
    "print(at_least_one_unigram(query,inverted_index))\n",
    "print(ngrams(query,inverted_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that our collection is:\n",
    "```python\n",
    "docs = ['The quick brown fox jumps over the lazy dog',\n",
    "        'The brown quick fox jumps over the lazy dog',\n",
    "        'Luke is the mechanical and electrical engineer of the group',\n",
    "        'Instead of buying a new engine, buy a new car',\n",
    "        'An engineer may design car engines of all sorts',\n",
    "        'Engineers use logic, but not necessarily imagination',\n",
    "        'Logic will take you from A to Z, imagination will take you everywhere.',\n",
    "        'Continuous effort, not strength or intelligence, is the key to \\\n",
    "        unlocking our potential. And curiosity.',\n",
    "        'It’s OK to have your eggs in one basket as long as you control what \\\n",
    "        happens to that basket.'\n",
    "        ]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TF-IDF\n",
    "Let's now compute the IDF coefficients for our words. These coefficients are a measure of term specificity. They are computed as:\n",
    "\n",
    "\\begin{equation}\n",
    "idf(t,D) = \\log \\big( \\frac{m}{1+df(t)} \\big)\n",
    "\\end{equation}\n",
    "\n",
    "where $df(t)$ is the number of documents in the collection $D$ that contain $t$ and $m$ is the size of $D$. The IDF coefficients are usually then plugged in the equation below to compute feature vetors for the documents:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathrm{weight}(t,d,D) = tf(t,d) \\times idf(t,D)\n",
    "\\end{equation}\n",
    "\n",
    "Where $tf(t,d)$ is the number of times term $t$ appears in document $d$. Here, the documents are represented as **bags of words**.\n",
    "\n",
    "\n",
    "The intuition behind bag-of-words + TF-IDF coefficients is that *frequent words in a document are representative of that document as long as they are not also very frequent at the corpus level*. \n",
    "\n",
    "In practice, the TF-IDF vectors can be used (1) in querying, to compute relevance scores for the documents, and (2) as features for (un)supervised tasks such as document clustering or classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('the', 0.25527), (',', 0.25527), ('.', 0.35218), ('of', 0.35218), ('to', 0.35218), ('lazy', 0.47712), ('car', 0.47712), ('is', 0.47712), ('fox', 0.47712), ('not', 0.47712), ('engineer', 0.47712), ('imagination', 0.47712), ('and', 0.47712), ('quick', 0.47712), ('over', 0.47712), ('dog', 0.47712), ('brown', 0.47712), ('jumps', 0.47712), ('a', 0.47712), ('logic', 0.47712), ('you', 0.47712), ('an', 0.65321), ('may', 0.65321), ('necessarily', 0.65321), ('new', 0.65321), ('curiosity', 0.65321), ('effort', 0.65321), ('everywhere', 0.65321), ('in', 0.65321), ('buying', 0.65321), ('strength', 0.65321), ('use', 0.65321), ('as', 0.65321), ('engines', 0.65321), ('engine', 0.65321), ('group', 0.65321), ('z', 0.65321), ('or', 0.65321), ('engineers', 0.65321), ('take', 0.65321), ('from', 0.65321), ('your', 0.65321), ('continuous', 0.65321), ('it', 0.65321), ('long', 0.65321), ('that', 0.65321), ('but', 0.65321), ('happens', 0.65321), ('electrical', 0.65321), ('sorts', 0.65321), ('buy', 0.65321), ('have', 0.65321), ('’', 0.65321), ('eggs', 0.65321), ('one', 0.65321), ('potential', 0.65321), ('will', 0.65321), ('mechanical', 0.65321), ('our', 0.65321), ('basket', 0.65321), ('s', 0.65321), ('all', 0.65321), ('ok', 0.65321), ('control', 0.65321), ('what', 0.65321), ('unlocking', 0.65321), ('instead', 0.65321), ('key', 0.65321), ('intelligence', 0.65321), ('luke', 0.65321), ('design', 0.65321)]\n"
     ]
    }
   ],
   "source": [
    "all_unique_terms = list(set([elt for sublist in cleaned_docs for elt in sublist])) # flatten docs\n",
    "\n",
    "terms_by_doc_sets = [set(elt) for elt in cleaned_docs]\n",
    "n_doc = len(cleaned_docs)\n",
    "idfs = dict(zip(all_unique_terms,[0]*len(all_unique_terms)))\n",
    "dfs = dict(zip(all_unique_terms,[0]*len(all_unique_terms)))\n",
    "\n",
    "for unique_term in list(idfs.keys()):\n",
    "    df = sum([unique_term in terms for terms in terms_by_doc_sets]) # nb of docs in which 'unique_term' appears\n",
    "    dfs[unique_term] = df\n",
    "    idfs[unique_term] = round(math.log10(n_doc/(1+df)),5)\n",
    "\n",
    "sorted_idfs = sorted(idfs.items(), key=operator.itemgetter(1), reverse=False)\n",
    "\n",
    "# terms sorted in increasing order of specificity\n",
    "print(sorted_idfs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course, we have a very small toy collection, so the picture is quite blurry here. But we can still see that the stopwords and punctuation marks have the lowest IDF scores (lowest specificity), which makes sense. Note that in a real corpus, there will be orders of magnitude differences between the IDF scores of the stopwords and that of the most specific words. Stopwords have very high TF values in all documents, but since their IDF scores are very low, their resulting final TF-IDF scores will be small."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
